Added cholesky inverse operation to all the backends #21554
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @buildwithsuhana, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request expands Keras's linear algebra capabilities by adding a cholesky_inverse function. This addition provides a robust and efficient method for matrix inversion, crucial for various numerical computations, and ensures consistent behavior across different backend frameworks.
Highlights
- New Linear Algebra Operation: I've introduced a
cholesky_inverseoperation, which provides a numerically stable and efficient way to compute the inverse of symmetric positive-definite matrices. This method first performs a Cholesky decomposition and then uses triangular solves to find the inverse. - Backend Implementations: The
cholesky_inversefunction has been implemented across multiple Keras backends: JAX, NumPy, TensorFlow, and PyTorch. For OpenVINO, it's currently marked as not supported. - Keras Ops Integration and Testing: The new operation is integrated into
keras.ops.linalg, supporting symbolic tensors and including input validation for 2D and square matrices. Corresponding unit tests have been added tolinalg_test.pyto verify its behavior and error handling.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces the cholesky_inverse operation across various Keras backends. The implementation is a welcome addition, providing a more efficient and numerically stable method for inverting symmetric positive-definite matrices compared to a general-purpose inverse. My review has identified a few issues, including critical bugs in the JAX and NumPy backend implementations related to handling batched inputs, which could lead to incorrect results. I've also noted a lack of correctness tests and a missing docstring, which are important for maintainability and usability. I've provided specific suggestions to address these points. Once these issues are resolved, this will be a solid contribution to the Keras API.
There was a problem hiding this comment.
Thanks for the PR, quick review! Let's run code formatting and linting. Here are the instructions: #21554 (comment).
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
|
@buildwithsuhana - can you set up pre-commit: https://github.com/keras-team/keras/blob/master/CONTRIBUTING.md#generating-public-api-and-formatting-the-code? This will ensure that your changes are formatted, and the symbols you've added are exported |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #21554 +/- ##
==========================================
+ Coverage 82.73% 82.75% +0.02%
==========================================
Files 567 567
Lines 56456 56528 +72
Branches 8823 8823
==========================================
+ Hits 46710 46782 +72
- Misses 7582 7585 +3
+ Partials 2164 2161 -3
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
google-ml-butler
bot
added
kokoro:force-run
ready to pull
|
@hertschuh - could you please drop an approval here? I don't have write access. |
google-ml-butler
bot
removed
the
ready to pull
google-ml-butler
bot
added
kokoro:force-run
ready to pull
github-project-automation
bot
moved this from Assigned Reviewer
to Approved by Reviewer
in PR Queue
github-project-automation
bot
moved this from Approved by Reviewer
to Merged
in PR Queue
google-ml-butler
bot
removed
awaiting review
ready to pull
The Cholesky inverse is a fast and numerically stable algorithm for calculating the inverse of a special type of matrix: one that is symmetric and positive-definite.
Instead of calculating the inverse directly, this method first breaks the matrix down into simpler components using Cholesky decomposition. This approach is generally more efficient and less prone to floating-point errors than a general-purpose inversion method.
Its numerical stability makes it a key component in advanced machine learning algorithms like GPTQ quantisation, where it is used to efficiently compute the inverse of Hessian matrices for layer-wise weight compression.